package crypto

import (
	"crypto/cipher"
	"crypto/elliptic"
	"crypto/rand"
	"errors"
	"fmt"

	"github.com/tjfoc/gmsm/sm2"
	"github.com/tjfoc/gmsm/sm3"
	"github.com/tjfoc/gmsm/sm4"
	"github.com/tjfoc/gmsm/x509"
)

type sm2KeyGenerator struct {
	curve elliptic.Curve
}

//GetCurve get curve
func (kg *sm2KeyGenerator) GetCurve() elliptic.Curve {
	return kg.curve
}

type sm2PrivateKey struct {
	privKey *sm2.PrivateKey
}

// Bytes converts this key to its byte representation,
// if this operation is allowed.
func (k *sm2PrivateKey) Bytes() (raw []byte) {
	priv := k.privKey
	raw = paddedBigBytes(priv.D, priv.Params().BitSize/8)
	return
}

// SKI returns the subject key identifier of this key.
func (k *sm2PrivateKey) SKI() (ski []byte) {
	if k.privKey == nil {
		return nil
	}

	// Marshall the public key
	raw := elliptic.Marshal(k.privKey.Curve, k.privKey.PublicKey.X, k.privKey.PublicKey.Y)

	// Hash it
	hash := sm3.New()
	hash.Write(raw)
	return hash.Sum(nil)
}

// IsSymmetric returns true if this key is a symmetric key,
// false if this key is asymmetric
func (k *sm2PrivateKey) IsSymmetric() bool {
	return false
}

// IsPrivate returns true if this key is a private key,
// false otherwise.
func (k *sm2PrivateKey) IsPrivate() bool {
	return true
}

// PublicKey returns the corresponding public key part of an asymmetric public/private key pair.
// This method returns an error in symmetric key schemes.
func (k *sm2PrivateKey) PublicKey() PublicKey {
	return &sm2PublicKey{&k.privKey.PublicKey}
}

//Sign sign
func (k *sm2PrivateKey) Sign(digest []byte) ([]byte, error) {
	r, s, err := sm2.Sm2Sign(k.privKey, digest, nil, rand.Reader)
	if err != nil {
		return nil, err
	}
	return marshalECDSASignature(r, s)
}

//GetSignatureAlgorithm get signature algorithm
func (k *sm2PrivateKey) GetSignatureAlgorithm() x509.SignatureAlgorithm {
	//switch options.HashFamily {
	//case "SM3":
	//	return x509.SM2WithSM3
	//case "SHA1":
	//	return x509.SM2WithSHA1
	//case "SHA","SHA256":
	//	return x509.SM2WithSHA256
	//}
	return x509.SM2WithSM3
}

type sm2PublicKey struct {
	pubKey *sm2.PublicKey
}

// Bytes converts this key to its byte representation,
// if this operation is allowed.
func (k *sm2PublicKey) Bytes() (raw []byte) {
	pub := k.pubKey
	if pub == nil || pub.X == nil || pub.Y == nil {
		return nil
	}
	raw = elliptic.Marshal(pub.Curve, pub.X, pub.Y)
	return
}

// SKI returns the subject key identifier of this key.
func (k *sm2PublicKey) SKI() (ski []byte) {
	if k.pubKey == nil {
		return nil
	}

	// Marshall the public key
	raw := elliptic.Marshal(k.pubKey.Curve, k.pubKey.X, k.pubKey.Y)

	// Hash it
	hash := sm3.New()
	hash.Write(raw)
	return hash.Sum(nil)
}

// IsSymmetric returns true if this key is a symmetric key,
// false if this key is asymmetric
func (k *sm2PublicKey) IsSymmetric() bool {
	return false
}

// IsPrivate returns true if this key is a private key,
// false otherwise.
func (k *sm2PublicKey) IsPrivate() bool {
	return false
}

//Verify verify signature
func (k *sm2PublicKey) Verify(message []byte, signature []byte) (bool, error) {
	r, s, err := unmarshalECDSASignature(signature)
	if err != nil {
		return false, fmt.Errorf("Failed unmashalling signature [%s]", err)
	}

	return sm2.Sm2Verify(k.pubKey, message, nil, r, s), nil
}

//Encrypt Encrypt
func (k *sm2PublicKey) Encrypt(data []byte) ([]byte, error) {
	return sm2.Encrypt(k.pubKey, data, rand.Reader)
}

//Decrypt Decrypt
func (k *sm2PrivateKey) Decrypt(data []byte) ([]byte, error) {
	return sm2.Decrypt(k.privKey, data)
}

//KeyGen generate key pair
func (kg *sm2KeyGenerator) KeyGen() (k Key, err error) {
	privKey, err := sm2.GenerateKey(nil)
	if err != nil {
		return nil, fmt.Errorf("Failed generating SM2 key for [%v]: [%s]", kg.curve, err)
	}

	return &sm2PrivateKey{privKey}, nil
}

type sm4KeyGenerator struct {
}
type sm4PrivateKey struct {
	privKey    []byte
	exportable bool
}

// Bytes converts this key to its byte representation,
// if this operation is allowed.
func (k *sm4PrivateKey) Bytes() (raw []byte) {
	if k.exportable {
		raw = k.privKey
	} else {
		raw = nil
	}
	return
}

// SKI returns the subject key identifier of this key.
func (k *sm4PrivateKey) SKI() (ski []byte) {
	hash := sm3.New()
	hash.Write([]byte{0x01})
	hash.Write(k.privKey)
	return hash.Sum(nil)
}

// IsSymmetric returns true if this key is a symmetric key,
// false if this key is asymmetric
func (k *sm4PrivateKey) IsSymmetric() bool {
	return true
}

// IsPrivate returns true if this key is a private key,
// false otherwise.
func (k *sm4PrivateKey) IsPrivate() bool {
	return true
}

//Encrypt encrypt
func (k *sm4PrivateKey) Encrypt(data []byte) ([]byte, error) {
	ciphertext := options.Padding.Padding(data)
	if len(ciphertext)%sm4.BlockSize != 0 {
		return nil, errors.New("Invalid plaintext. It must be a multiple of the block size")
	}
	block, err := sm4.NewCipher(k.privKey)
	if err != nil {
		return nil, err
	}
	hash, _ := Hash(k.privKey)
	iv := hash[:sm4.BlockSize]

	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext, ciphertext)

	return ciphertext, nil
}

//Decrypt decrypt
func (k *sm4PrivateKey) Decrypt(src []byte) ([]byte, error) {
	block, err := sm4.NewCipher(k.privKey)
	if err != nil {
		return nil, err
	}

	if len(src) < sm4.BlockSize {
		return nil, errors.New("Invalid ciphertext. It must be a multiple of the block size")
	}
	hash, _ := Hash(k.privKey)
	iv := hash[:sm4.BlockSize]

	if len(src)%sm4.BlockSize != 0 {
		return nil, errors.New("Invalid ciphertext. It must be a multiple of the block size")
	}

	mode := cipher.NewCBCDecrypter(block, iv)

	mode.CryptBlocks(src, src)
	return options.Padding.UnPadding(src)

}

//KeyGen generate key pair
func (kg *sm4KeyGenerator) KeyGen() (k Key, err error) {
	lowLevelKey, err := RandomBytes(16)
	if err != nil {
		return nil, fmt.Errorf("Failed generating SM4 key [%s]", err)
	}
	return &sm4PrivateKey{lowLevelKey, true}, nil
}
